'''
furnet in PyTorch.
熔炉温控关系网络
'''
import torch.nn as nn
import torch.nn.functional as F
from .init_utils import weights_init


class FurNet(nn.Module):
    def __init__(self, inx=8, out=6):
        super(FurNet, self).__init__()
        self.fc1 = nn.Linear(inx, 64)
        self.fc2 = nn.Linear(64, 128)
        self.fc3 = nn.Linear(128, out)
        # self.ttt = nn.Linear(inx, out)
        self.apply(weights_init)

    def forward(self, x):
        out = F.relu(self.fc1(x))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        # out = self.ttt(x)
        return out


def furnet(dataset):
    if dataset == 'furnace86':
        return FurNet(8, 6)
    return FurNet()

